import { zodResolver } from "@hookform/resolvers/zod";
import {
  BoxesIcon,
  ChevronRightIcon,
  LoaderCircleIcon,
  Trash2Icon,
  WrenchIcon,
} from "lucide-react";
import * as React from "react";
import { useForm } from "react-hook-form";
import { toast } from "sonner";
import { z } from "zod";
import {
  apiAddConnection,
  apiAddNode,
  apiGetDefaultProperty,
  apiRemoveNode,
  apiReplaceNodeModule,
  isLLM,
} from "@/common";
import {
  AddonDef,
  Destination,
  type Graph,
  ProtocolLabel as GraphConnProtocol,
  GraphEditor,
  ProtocolLabel,
} from "@/common/graph";
import { useAppSelector, useGraphs } from "@/common/hooks";
import {
  compatibleTools,
  ModuleRegistry,
  ModuleTypeLabels,
  moduleRegistry,
  toolModuleRegistry,
} from "@/common/moduleConfig";
import { Button, buttonVariants } from "@/components/ui/button";
import {
  Form,
  FormControl,
  FormDescription,
  FormField,
  FormItem,
  FormLabel,
  FormMessage,
} from "@/components/ui/form";
import {
  Select,
  SelectContent,
  SelectGroup,
  SelectItem,
  SelectLabel,
  SelectTrigger,
  SelectValue,
} from "@/components/ui/select";
import {
  Sheet,
  SheetContent,
  SheetDescription,
  SheetHeader,
  SheetTitle,
  SheetTrigger,
} from "@/components/ui/sheet";
import { cn } from "@/lib/utils";
import { fetchGraphDetails } from "@/store/reducers/global";
import {
  DropdownMenu,
  DropdownMenuContent,
  DropdownMenuItem,
  DropdownMenuPortal,
  DropdownMenuSub,
  DropdownMenuSubContent,
  DropdownMenuSubTrigger,
  DropdownMenuTrigger,
} from "../ui/dropdown";

export function RemoteModuleCfgSheet() {
  const addonModules = useAppSelector((state) => state.global.addonModules);
  const {
    getGraphNodeAddonByName,
    selectedGraph,
    updateGraph,
    installedAndRegisteredModulesMap,
    installedAndRegisteredToolModules,
  } = useGraphs();

  const metadata = React.useMemo(() => {
    const dynamicMetadata: Record<
      string,
      { type: string; options: { value: string; label: string }[] }
    > = {};

    if (selectedGraph) {
      Object.keys(installedAndRegisteredModulesMap).forEach((key) => {
        const moduleTypeKey = key as ModuleRegistry.NonToolModuleType;

        // Check if the current graph has a node whose name contains the ModuleType
        const hasMatchingNode = selectedGraph.nodes.some((node) =>
          node.name.includes(moduleTypeKey)
        );

        if (hasMatchingNode) {
          dynamicMetadata[moduleTypeKey] = {
            type: "string",
            options: installedAndRegisteredModulesMap[moduleTypeKey].map(
              (module) => ({
                value: module.name,
                label: module.label,
              })
            ),
          };
        }
      });
    }

    return dynamicMetadata;
  }, [installedAndRegisteredModulesMap, selectedGraph]);

  const initialData = React.useMemo(() => {
    const dynamicInitialData: Record<string, string | null | undefined> = {};

    if (selectedGraph) {
      Object.keys(installedAndRegisteredModulesMap).forEach((key) => {
        const moduleTypeKey = key as ModuleRegistry.ModuleType;

        // Check if the current graph has a node whose name contains the ModuleType
        const hasMatchingNode = selectedGraph.nodes.some((node) =>
          node.name.includes(moduleTypeKey)
        );

        if (hasMatchingNode) {
          dynamicInitialData[moduleTypeKey] =
            getGraphNodeAddonByName(moduleTypeKey)?.addon;
        }
      });
    }

    return dynamicInitialData;
  }, [
    installedAndRegisteredModulesMap,
    selectedGraph,
    getGraphNodeAddonByName,
  ]);

  return (
    <Sheet>
      <SheetTrigger
        className={cn(
          buttonVariants({ variant: "outline", size: "icon" }),
          "bg-transparent"
        )}
      >
        <BoxesIcon />
      </SheetTrigger>
      <SheetContent className="w-[400px] overflow-y-auto sm:w-[540px]">
        <SheetHeader>
          <SheetTitle>Module Picker</SheetTitle>
          <SheetDescription>
            You can adjust STT/TTS/LLM/LLMv2v extension modules here, the values
            will be written into property.json file when you save.
          </SheetDescription>
        </SheetHeader>

        <div className="my-4">
          <GraphModuleCfgForm
            initialData={initialData}
            metadata={metadata}
            onUpdate={async (data, tools) => {
              try {
                // Clone the selectedGraph to avoid mutating the original graph
                const selectedGraphCopy: Graph = JSON.parse(
                  JSON.stringify(selectedGraph)
                );
                const nodes = selectedGraphCopy.nodes;
                let needUpdate = false;
                let enableRTCVideoSubscribe = false;

                // Retrieve the agora_rtc node
                const agoraRtcNode = GraphEditor.findNode(
                  selectedGraphCopy,
                  "agora_rtc"
                );
                if (!agoraRtcNode) {
                  toast.error("agora_rtc node not found in the graph");
                  return;
                }

                // Update graph nodes with selected modules
                Object.entries(data).forEach(([key, value]) => {
                  const node = nodes.find((n) => n.name === key);
                  if (node && value && node.addon !== value) {
                    node.addon = value;
                    node.property = addonModules.find(
                      (module) => module.name === value
                    )?.defaultProperty;
                    needUpdate = true;
                  }
                });

                const reasoningNodesWithVisualSupport =
                  GraphEditor.findNodeByPredicate(selectedGraphCopy, (node) => {
                    const module = moduleRegistry[node.addon];
                    if (!module) {
                      return false;
                    }

                    if (module.type === ModuleRegistry.ModuleType.LLM) {
                      const llmModule = module as ModuleRegistry.LLMModule;
                      return (
                        isLLM(node.name) &&
                        llmModule.options.inputModalities.includes(
                          ModuleRegistry.Modalities.Video
                        )
                      );
                    }

                    if (module.type === ModuleRegistry.ModuleType.V2V) {
                      const v2vModule = module as ModuleRegistry.V2VModule;
                      return (
                        isLLM(node.name) &&
                        v2vModule.options.inputModalities.includes(
                          ModuleRegistry.Modalities.Video
                        )
                      );
                    }

                    return false;
                  });

                if (reasoningNodesWithVisualSupport) {
                  GraphEditor.addOrUpdateConnection(
                    selectedGraphCopy,
                    `${agoraRtcNode.name}`,
                    `${reasoningNodesWithVisualSupport.name}`,
                    ProtocolLabel.VIDEO_FRAME,
                    "video_frame"
                  );
                  enableRTCVideoSubscribe = true;
                }

                // Identify removed tools and process them
                const currentToolsInGraph = nodes
                  .filter((node) =>
                    installedAndRegisteredToolModules
                      .map((module) => module.name)
                      .includes(node.addon)
                  )
                  .map((node) => node.addon);

                const removedTools = currentToolsInGraph.filter(
                  (tool) => !tools.includes(tool)
                );
                removedTools.forEach((tool) => {
                  GraphEditor.removeNodeAndConnections(selectedGraphCopy, tool);
                  needUpdate = true;
                });

                // Process tool modules
                if (tools.length > 0) {
                  if (!enableRTCVideoSubscribe) {
                    enableRTCVideoSubscribe = tools.some((tool) =>
                      tool.includes("vision")
                    );
                  }
                  tools.forEach((tool) => {
                    if (!currentToolsInGraph.includes(tool)) {
                      const toolModule = addonModules.find(
                        (module) => module.name === tool
                      );

                      if (!toolModule) {
                        toast.error(`Module ${tool} not found`);
                        return;
                      }

                      const toolNode = GraphEditor.addNode(
                        selectedGraphCopy,
                        tool,
                        tool,
                        "default",
                        toolModule.defaultProperty
                      );

                      // Create or update connections
                      const llmNode = GraphEditor.findNodeByPredicate(
                        selectedGraphCopy,
                        (node) => isLLM(node.name)
                      );
                      if (llmNode) {
                        GraphEditor.linkTool(
                          selectedGraphCopy,
                          llmNode,
                          toolNode,
                          toolModuleRegistry[tool]
                        );
                      }
                    }
                  });
                  needUpdate = true;
                }

                GraphEditor.enableRTCVideoSubscribe(
                  selectedGraphCopy,
                  enableRTCVideoSubscribe
                );

                // Perform the update if changes are detected
                if (needUpdate) {
                  await updateGraph(selectedGraph, selectedGraphCopy);
                  toast.success("Modules updated", {
                    description: `Graph: ${selectedGraphCopy.graph_id}`,
                  });
                }
              } catch (e: any) {
                toast.error(`Failed to update modules: ${e}`);
              }
            }}
          />
        </div>
      </SheetContent>
    </Sheet>
  );
}

const GraphModuleCfgForm = ({
  initialData,
  metadata,
  onUpdate,
}: {
  initialData: Record<string, string | null | undefined>;
  metadata: Record<
    string,
    { type: string; options: { value: string; label: string }[] }
  >;
  onUpdate: (data: Record<string, string | null>, tools: string[]) => void;
}) => {
  const formSchema = z.record(z.string(), z.string().nullable());
  const form = useForm<z.infer<typeof formSchema>>({
    resolver: zodResolver(formSchema),
    defaultValues: initialData,
  });
  const { selectedGraph, installedAndRegisteredToolModules } = useGraphs();
  const { watch } = form;

  // Watch for changes in "llm" and "v2v" fields
  const llmValue = watch("llm");
  const v2vValue = watch("v2v");
  const toolModules = React.useMemo(() => {
    // Step 1: Get installed and registered tool modules
    const allToolModules = installedAndRegisteredToolModules || [];

    // Step 2: Determine the active module based on form values
    const activeModule = llmValue || v2vValue;

    // Step 3: Get compatible tools for the active module
    if (activeModule) {
      const compatibleToolNames = compatibleTools[activeModule] || [];
      return allToolModules.filter((module) =>
        compatibleToolNames.includes(module.name)
      );
    }

    // If no LLM or V2V module is selected, return all tool modules
    return [];
  }, [installedAndRegisteredToolModules, selectedGraph, llmValue, v2vValue]);

  const onSubmit = (data: z.infer<typeof formSchema>) => {
    onUpdate(data, selectedTools);
  };

  const [selectedTools, setSelectedTools] = React.useState<string[]>([]);

  // Synchronize selectedTools with selectedGraph and toolModules
  React.useEffect(() => {
    const toolNames = toolModules.map((module) => module.name);
    const graphToolAddons =
      selectedGraph?.nodes
        .filter((node) => toolNames.includes(node.addon))
        .map((node) => node.addon) || [];
    setSelectedTools(graphToolAddons);
  }, [toolModules, selectedGraph]);

  // Desired field order
  const fieldOrder: ModuleRegistry.NonToolModuleType[] = [
    ModuleRegistry.ModuleType.STT,
    ModuleRegistry.ModuleType.LLM,
    ModuleRegistry.ModuleType.V2V,
    ModuleRegistry.ModuleType.TTS,
  ];
  return (
    <Form {...form}>
      <form onSubmit={form.handleSubmit(onSubmit)} className="space-y-6">
        {fieldOrder.map(
          (key) =>
            metadata[key] && ( // Check if the field exists in metadata
              <div key={key} className="space-y-2">
                <FormField
                  control={form.control}
                  name={key}
                  render={({ field }) => (
                    <FormItem>
                      <FormLabel>
                        <div className="flex items-center justify-between">
                          <div className="py-3">{ModuleTypeLabels[key]}</div>
                          {isLLM(key) && (
                            <DropdownMenu>
                              <DropdownMenuTrigger
                                className={cn(
                                  buttonVariants({
                                    variant: "outline",
                                    size: "icon",
                                  }),
                                  "bg-transparent"
                                )}
                              >
                                <WrenchIcon />
                              </DropdownMenuTrigger>
                              <DropdownMenuContent>
                                <DropdownMenuSub>
                                  <DropdownMenuSubTrigger
                                    icon={<ChevronRightIcon size={15} />}
                                    className="flex justify-between"
                                  >
                                    Add Tools
                                  </DropdownMenuSubTrigger>
                                  <DropdownMenuPortal>
                                    <DropdownMenuSubContent
                                      className="DropdownMenuSubContent"
                                      sideOffset={2}
                                      alignOffset={-5}
                                    >
                                      {toolModules.length > 0 ? (
                                        toolModules.map((module) => (
                                          <DropdownMenuItem
                                            key={module.name}
                                            disabled={selectedTools.includes(
                                              module.name
                                            )} // Disable if the tool is already selected
                                            onClick={() => {
                                              if (
                                                !selectedTools.includes(
                                                  module.name
                                                )
                                              ) {
                                                setSelectedTools((prev) => [
                                                  ...prev,
                                                  module.name,
                                                ]);
                                              }
                                            }}
                                          >
                                            {module.name}
                                          </DropdownMenuItem>
                                        ))
                                      ) : (
                                        <DropdownMenuItem disabled>
                                          No compatible tools
                                        </DropdownMenuItem>
                                      )}
                                    </DropdownMenuSubContent>
                                  </DropdownMenuPortal>
                                </DropdownMenuSub>
                              </DropdownMenuContent>
                            </DropdownMenu>
                          )}
                        </div>
                      </FormLabel>
                      <FormControl>
                        <Select
                          value={field.value ?? ""}
                          onValueChange={field.onChange}
                        >
                          <SelectTrigger>
                            <SelectValue
                              placeholder={`Select a ${key.toUpperCase()} option`}
                            />
                          </SelectTrigger>
                          <SelectContent>
                            {metadata[key].options.map((option) => (
                              <SelectItem
                                key={option.value}
                                value={option.value}
                              >
                                {option.label}
                              </SelectItem>
                            ))}
                          </SelectContent>
                        </Select>
                      </FormControl>
                    </FormItem>
                  )}
                />
                {isLLM(key) && selectedTools.length > 0 && (
                  <div className="mt-2">
                    {selectedTools.map((tool) => (
                      <div
                        key={tool}
                        className="flex items-center justify-between rounded-md px-3 py-2"
                      >
                        <span className="text-sm">{tool}</span>
                        <div
                          className={cn(
                            buttonVariants({
                              variant: "outline",
                              size: "icon",
                            }),
                            "bg-transparent",
                            "ml-2",
                            "cursor-pointer"
                          )}
                          onClick={() => {
                            setSelectedTools((prev) =>
                              prev.filter((t) => t !== tool)
                            );
                          }} // Delete action
                        >
                          <Trash2Icon />
                        </div>
                      </div>
                    ))}
                  </div>
                )}
              </div>
            )
        )}

        <Button type="submit" disabled={form.formState.isSubmitting}>
          {form.formState.isSubmitting ? (
            <>
              <LoaderCircleIcon className="h-4 w-4 animate-spin" />
              <span>Saving...</span>
            </>
          ) : (
            "Save changes"
          )}
        </Button>
      </form>
    </Form>
  );
};
